# -*- coding: utf-8 -*-
# Author: lkm
# date: 2022/3/9 21:16

from pandas import read_csv, set_option
from lifelines import KaplanMeierFitter
from lifelines.utils import median_survival_times
from lifelines import CoxPHFitter
from matplotlib import pyplot as plt
from toad.detector import detect

set_option("display.width", 10000)
set_option("display.max_rows", None)
set_option("display.max_columns", None)


class SurvivalAnalysis(object):
    def __init__(self, file):
        data = read_csv(file)

        print(data.head())

        kmf = KaplanMeierFitter()
        kmf.fit(data['tenure'], event_observed=data['Churn_Yes'])

        kmf.plot_survival_function()

        median_ = kmf.median_survival_time_
        median_confidence_interval_ = median_survival_times(kmf.confidence_interval_)
        # print(median_confidence_interval_)
        # plt.show()

        cph = CoxPHFitter()
        dt = data[["tenure", "Churn_Yes", "SeniorCitizen", "MonthlyCharges", "TotalCharges"]]
        cph.fit(dt, 'tenure', event_col='Churn_Yes')
        print()
        print(cph.print_summary())


if __name__ == "__main__":
    filename = "data/precessed_feature.csv"
    SurvivalAnalysis(filename)
